import logging
import time
import zmq
from zmq.eventloop.ioloop import PeriodicCallback
from zmq.eventloop.zmqstream import ZMQStream
from bstar import BinaryStar
from kvsimple import KVMsg
from zhelpers import dump


class Route:
    def __init__(self, socket, identity, subtree):
        self.socket = socket
        self.identity = identity
        self.subtree = subtree


def send_single(key, kvmsg, route):
    if kvmsg.key.startswith(route.subtree):
        route.socket.send(route.identity, zmq.SNDMORE)
        kvmsg.send(route.socket)


class CloneServer(object):
    ctx = None
    kvmap = None
    bstar = None
    sequence = 0
    port = None
    peer = None
    publisher = None
    collector = None
    subscriber = None
    pending = None
    primary = False
    master = False
    slave = False

    def __init__(self, primary=True, ports=(5556, 5566)):
        self.primary = primary
        if primary:
            self.port, self.peer = ports
            frontend = "tcp://*:5003"
            backend = "tcp://localhost:5004"
            self.kvmap = {}
        else:
            self.peer, self.port = ports
            frontend = "tcp://*:5004"
            backend = "tcp://localhost:5003"

        self.ctx = zmq.Context.instance()
        self.pending = []
        self.bstar = BinaryStar(primary, frontend, backend)
        self.bstar.register_voter("tcp://*:%i" % self.port, zmq.ROUTER, self.handle_snapshot)

        self.bstar.master_callback = self.become_master
        self.bstar.slave_callback = self.become_slave

        self.publisher = self.ctx.socket(zmq.PUB)
        self.collector = self.ctx.socket(zmq.SUB)
        self.collector.setsockopt(zmq.SUBSCRIBE, b'')
        self.publisher.bind("tcp://*:%d" % (self.port + 1))
        self.collector.bind("tcp://*:%d" % (self.port + 2))

        self.subscriber = self.ctx.socket(zmq.SUB)
        self.subscriber.setsockopt(zmq.SUBSCRIBE, b'')
        self.subscriber.connect("tcp://localhost:%d" % (self.peer + 1))

        self.publisher = ZMQStream(self.publisher)
        self.subscriber = ZMQStream(self.subscriber)
        self.collector = ZMQStream(self.collector)

        self.collector.on_recv(self.handle_collect)
        self.flush_callback = PeriodicCallback(self.flush_ttl, 1000)
        self.hugz_callback = PeriodicCallback(self.send_hugz, 1000)

        logging.basicConfig(format="%(asctime)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)

    def start(self):
        self.flush_callback.start()
        self.hugz_callback.start()
        try:
            self.bstar.start()
        except KeyboardInterrupt:
            pass

    def handle_snapshot(self, socket, msg):
        if msg[1] != b"ICANHAZ?" or len(msg) != 3:
            logging.error("E: bad request, aborting")
            dump(msg)
            self.bstar.loop.stop()
            return
        identity, request = msg[:2]
        if len(msg) >= 3:
            subtree = msg[2].decode('utf-8')
            route = Route(socket, identity, subtree)
            for k, v in self.kvmap.items():
                send_single(k, v, route)

            logging.info("I: Sending state snapshot=%d" % self.sequence)
            socket.send(identity, zmq.SNDMORE)
            kvmsg = KVMsg(self.sequence)
            kvmsg.key = "KTHXBAI"
            kvmsg.body = subtree
            kvmsg.send(socket)

    def become_master(self):
        self.master = True
        self.slave = False
        self.subscriber.stop_on_recv()

        while self.pending:
            kvmsg = self.pending.pop(0)
            self.sequence += 1
            kvmsg.sequence = self.sequence
            kvmsg.store(self.kvmap)
            logging.info("I: publishing pending=%d", self.sequence)

    def become_slave(self):
        self.kvmap = None
        self.master = False
        self.slave = True
        self.subscriber.on_recv(self.handle_subscriber)

    def handle_collect(self, msg):
        kvmsg = KVMsg.from_msg(msg)
        if self.master:
            self.sequence += 1
            kvmsg.sequence = self.sequence
            kvmsg.send(self.publisher)
            ttl = float(kvmsg.get(b'ttl', 0))
            if ttl:
                kvmsg[b'ttl'] = b'%f' % (time.time() + ttl)
            kvmsg.store(self.kvmap)
            # logging.info("I: publishing update=%d %s" % (self.sequence, kvmsg.body))
        else:
            if not self.was_pending(kvmsg):
                self.pending.append(kvmsg)

    def was_pending(self, kvmsg):
        found = False
        for idx, held in enumerate(self.pending):
            if held.uuid == kvmsg.uuid:
                found = True
                break
        if found:
            self.pending.pop(idx)
        return found

    def handle_subscriber(self, msg):
        if self.master:
            logging.warn("received subscriber message, but we are master %s" % msg)
            return
        if self.kvmap is None:
            self.kvmap = {}
            snapshot = self.ctx.socket(zmq.DEALER)
            snapshot.linger = 0
            snapshot.connect("tcp://localhost:%i" % self.peer)
            logging.info("I: asking for snapshot from: tcp://localhost:%d" % self.peer)
            snapshot.send_multipart([b"ICANHAZ?", b""])
            while True:
                try:
                    kvmsg = KVMsg.recv(snapshot)
                except KeyboardInterrupt:
                    self.bstar.loop.stop()
                    return
                if kvmsg.key == b"KTHXBAI":
                    self.sequence = kvmsg.sequence
                    break
                kvmsg.store(self.kvmap)
            logging.info("I: received snapshot=%d" % self.sequence)

        kvmsg = KVMsg.from_msg(msg)
        if kvmsg.key != b"HUGZ":
            if not self.was_pending(kvmsg):
                self.pending.append(kvmsg)

            if kvmsg.sequence > self.sequence:
                self.sequence = kvmsg.sequence
                kvmsg.store(self.kvmap)
                logging.info("I: received update=%d" % self.sequence)

    def flush_ttl(self):
        if self.kvmap:
            for key, kvmsg in list(self.kvmap.items()):
                self.flush_single(kvmsg)

    def flush_single(self, kvmsg):
        ttl = float(kvmsg.get(b'ttl', 0))
        if ttl and ttl <= time.time():
            kvmsg.body = ""
            self.sequence += 1
            kvmsg.sequence = self.sequence
            # logging.info("I: preparing to publish delete=%s", kvmsg.properties)
            kvmsg.send(self.publisher)
            del self.kvmap[kvmsg.key]
            # logging.info("I: publishing delete=%d" % self.sequence)

    def send_hugz(self):
        kvmsg = KVMsg(self.sequence)
        kvmsg.key = "HUGZ"
        kvmsg.body = ""
        # logging.info("I: sending hugz")
        kvmsg.send(self.publisher)


def main():
    import sys
    if '-p' in sys.argv:
        primary = True
    elif '-b' in sys.argv:
        primary = False
    else:
        print("Usage: clonesrv6.py {-p|-b}")
        sys.exit(1)

    clone = CloneServer(primary)
    clone.start()


if __name__ == "__main__":
    main()
